package test.demo.service;

import com.alibaba.fastjson.JSON;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import test.demo.dto.EventMessage;

import java.util.HashMap;
import java.util.Map;

@Service
public class EventStreamService {

    private final WebClient webClient;

    public EventStreamService() {
        this.webClient = WebClient.builder()
                .defaultHeader("Accept", "text/event-stream")
                .build();
    }

    /**
     * http://localhost:11434/api/generate
     * {
     *    model: "deepseek-r1:1.5b",
     *    prompt: '你好',
     *    stream: true,
     *  }
     * @param userInput
     * @param messageId
     * @return
     */
    public Flux<String> getEventStream(String userInput, String messageId) {
        Map<String, Object> bodyMap = new HashMap();
        bodyMap.put("model", "deepseek-r1:1.5b");
        bodyMap.put("prompt", userInput);
        bodyMap.put("stream", true);

        // 这个contentBuilder是把ollama 每次返回的增量字符，拼接成全量字符返回
        StringBuilder contentBuilder = new StringBuilder();
        return webClient.post()
                .uri("http://localhost:11434/api/generate")
                .header("Content-Type", "application/json") // 设置内容类型
                .bodyValue(JSON.toJSONString(bodyMap))
                .retrieve()
                .bodyToFlux(String.class)
                // 消息处理
                .flatMap(data -> processResponse(data,contentBuilder, messageId));
    }

    private Mono<String> processResponse(String data,StringBuilder contentBuilder, String messageId) {
        // 把ollma的消息转成java对象
        EventMessage response = JSON.parseObject(data, EventMessage.class);
        // 消息id返回给前端
        response.setMsg_id(messageId);
        // 拼接全量字符
        String content = response.getResponse();
        contentBuilder.append(content);
        response.setResponse(contentBuilder.toString());
        // 返回string
        return Mono.just(JSON.toJSONString(response));
    }
}
